from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage

class LoopState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]
    iteration_count: Annotated[int, "迭代次数"]
    max_iterations: Annotated[int, "最大迭代次数"]
    should_continue: Annotated[bool, "是否继续循环"]
    loop_history: Annotated[list, "循环历史"]

class LoopGraphDemo:
    """循环图演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建循环图"""
        
        workflow = StateGraph(LoopState)
        
        # 添加节点
        workflow.add_node("initialize_loop", self._initialize_loop)
        workflow.add_node("process_iteration", self._process_iteration)
        workflow.add_node("check_condition", self._check_condition)
        workflow.add_node("finalize_loop", self._finalize_loop)
        
        # 设置入口点
        workflow.set_entry_point("initialize_loop")
        
        # 设置条件边
        workflow.add_conditional_edges(
            "check_condition",
            self._route_condition,
            {
                "continue": "process_iteration",
                "end": "finalize_loop"
            }
        )
        
        # 设置其他边
        workflow.add_edge("initialize_loop", "process_iteration")
        workflow.add_edge("process_iteration", "check_condition")
        workflow.add_edge("finalize_loop", END)
        
        return workflow.compile()
    
    def _initialize_loop(self, state: LoopState) -> LoopState:
        """初始化循环"""
        messages = state.get("messages", [])
        if not messages:
            messages = [HumanMessage(content="开始循环处理")]
        
        max_iterations = state.get("max_iterations", 5)
        
        return {
            "messages": messages,
            "result": "循环已初始化",
            "iteration_count": 0,
            "max_iterations": max_iterations,
            "should_continue": True,
            "loop_history": ["循环初始化"]
        }
    
    def _process_iteration(self, state: LoopState) -> LoopState:
        """处理单次迭代"""
        iteration_count = state["iteration_count"] + 1
        max_iterations = state["max_iterations"]
        
        # 模拟迭代处理
        iteration_result = f"第{iteration_count}次迭代处理完成"
        
        # 更新循环历史
        loop_history = state["loop_history"] + [f"迭代{iteration_count}: {iteration_result}"]
        
        # 检查是否应该继续
        should_continue = iteration_count < max_iterations
        
        return {
            "messages": state["messages"],
            "result": iteration_result,
            "iteration_count": iteration_count,
            "max_iterations": max_iterations,
            "should_continue": should_continue,
            "loop_history": loop_history
        }
    
    def _check_condition(self, state: LoopState) -> LoopState:
        """检查循环条件"""
        should_continue = state["should_continue"]
        iteration_count = state["iteration_count"]
        
        # 更新循环历史
        loop_history = state["loop_history"] + [f"检查条件: 迭代{iteration_count}, 继续={should_continue}"]
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "iteration_count": iteration_count,
            "max_iterations": state["max_iterations"],
            "should_continue": should_continue,
            "loop_history": loop_history
        }
    
    def _route_condition(self, state: LoopState) -> str:
        """路由条件"""
        return "continue" if state["should_continue"] else "end"
    
    def _finalize_loop(self, state: LoopState) -> LoopState:
        """完成循环"""
        iteration_count = state["iteration_count"]
        loop_history = state["loop_history"]
        
        final_result = f"循环完成，共执行{iteration_count}次迭代"
        
        # 更新循环历史
        loop_history = loop_history + ["循环结束"]
        
        return {
            "messages": state["messages"] + [AIMessage(content=final_result)],
            "result": final_result,
            "iteration_count": iteration_count,
            "max_iterations": state["max_iterations"],
            "should_continue": False,
            "loop_history": loop_history
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行循环图"""
        max_iterations = input_data.get("max_iterations", 5)
        
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "开始循环处理"))],
            "result": "",
            "iteration_count": 0,
            "max_iterations": max_iterations,
            "should_continue": True,
            "loop_history": []
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "total_iterations": result["iteration_count"],
            "max_iterations": result["max_iterations"],
            "loop_history": result["loop_history"],
            "final_result": result["result"]
        }
